from aiModels.ollama_api import chat, DEFAULT_MODEL
from aiModels import SYSTEM_PROMPT
import re
import logging

# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# def get_deepseek_data(user_prompt, session=None):
#     if session is None:
#         logging.info(f"不存在的会话，创建一个新的会话")
#         session = [
#             {"role": "system", "content": SYSTEM_PROMPT},
#             {"role": "user", "content": user_prompt}
#         ]
#         session.append({"role": "assistant", "content": ""})
#     else:
#         logging.info(f"已存在会话，继续使用")
#         session.append({"role": "user", "content": user_prompt})
#         session.append({"role": "assistant", "content": ""})
#
#     # 使用流式输出调用 chat 函数
#     for chunk in chat(model=DEFAULT_MODEL, messages=session, stream=True):
#         # 去除<think>标签中的文本
#         if "<think>" in chunk:
#             chunk = re.sub(r'<think>.*?</think>', '', chunk, flags=re.DOTALL)
#         yield chunk
#
#     # 更新会话
#     full_content = ''.join([chunk for chunk in get_deepseek_data(user_prompt, session)])
#     session[-1]["content"] = full_content

def get_deepseek_data(user_prompt, session=None):
    if session is None:
        logging.info(f"不存在的会话，创建一个新的会话")
        session = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt}
        ]
        session.append({"role": "assistant", "content": ""})
    else:
        logging.info(f"已存在会话，继续使用")
        session.append({"role": "user", "content": user_prompt})
        session.append({"role": "assistant", "content": ""})
    full_content = ""
    # 调用 chat 函数获取完整响应
    for chunk in chat(model=DEFAULT_MODEL, messages=session, stream=True):
        # 去除<think>标签中的文本
        if "<think>" in chunk:
            chunk = re.sub(r'<think>.*?</think>', '', chunk, flags=re.DOTALL)
        full_content += chunk
    # 更新会话
    session[-1]["content"] = full_content
    return full_content

if __name__ == '__main__':
    session = None
    while True:
        user_input = input("用户：")
        print("助手：", end='', flush=True)
        full_content = ""
        for chunk in get_deepseek_data(user_prompt=user_input, session=session):
            print(chunk, end='', flush=True)
            full_content += chunk
        print()  # 打印换行符
        # 在这里可以使用 full_content 进行后续操作
        # 例如将其传递给下一个调用者的函数